"""嵌入API测试脚本"""
import os
import logging
from rag_system.embeddings import Embeddings

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def test_embeddings(backend: str = "deepseek"):
    """测试嵌入功能
    
    Args:
        backend: 要测试的后端，支持"deepseek"或"qwen"
    """
    try:
        logger.info(f"开始测试{backend}嵌入API...")
        
        # 初始化嵌入模型
        embeddings = Embeddings(backend=backend)
        
        # 测试单个文本嵌入
        text = "测试DeepSeek嵌入API"
        vector = embeddings.embed_text(text)
        logger.info(f"单个文本嵌入成功，向量长度: {len(vector)}")
        
        # 测试批量文本嵌入
        texts = ["测试文本1", "测试文本2", "测试文本3"]
        vectors = embeddings.embed_texts(texts)
        logger.info(f"批量文本嵌入成功，向量数量: {len(vectors)}")
        
        # 测试相似度计算
        sim = embeddings.similarity(vector, vectors[0])
        logger.info(f"相似度计算结果: {sim:.4f}")
        
        logger.info("所有测试通过！")
        return True
    except Exception as e:
        logger.error(f"测试失败: {str(e)}")
        return False

if __name__ == "__main__":
    # 测试DeepSeek后端
    try:
        test_embeddings("deepseek")
    except Exception as e:
        logger.error(f"DeepSeek测试失败: {str(e)}")
    
    # 测试通义千问后端(如果有配置)
    if os.getenv('QWEN_API_KEY'):
        try:
            test_embeddings("qwen")
        except Exception as e:
            logger.error(f"通义千问测试失败: {str(e)}")
    else:
        logger.warning("未配置QWEN_API_KEY，跳过通义千问测试")
